import pyjags
def dina_parameter_estimate(N,Q,I,K,simu_resp,chains=2,cores=4,burning=5000,iteration=5000,return_trace=False,return_full_vars=False):
    # hodina_pyjags
    jags_ho_dina_data = {
        "Y": simu_resp,
        "Q": Q,
        "N": N,
        "K": K,
        "I": I

    }
    if return_full_vars == True:
        jags_model_return_vars = ["p", "alpha_p", "alpha", "theta", "beta", "lambda_", "s", "g"]
    else :
        jags_model_return_vars = ["alpha"]

    jags_hodina_model_str =\
    """
    model{
        for (n in 1:N) {
            for (i in 1:I){
                for (k in 1:K){
                    w[n,i,k] <-(alpha[n,k]>=Q[i,k])
                }
                eta[n,i] <- prod(w[n,i,1:K])
                p[n,i] <- eta[n,i]*(1-s[i]-g[i]) + g[i]
                Y[n,i] ~ dbern(p[n,i])
            }
        }

        for (n in 1:N){
            for (k in 1:K){
                logit(alpha_p[n,k]) <- lambda_[k]*theta[n]+beta[k]
                alpha[n,k] ~ dbern(alpha_p[n,k]) 
            }
            theta[n] ~ dnorm(0,1)
        }

        for (k in 1:K){
            beta[k] ~ dnorm(0,0.25)
            lambda_[k] ~ dnorm(0,0.25) T(0,)
        } 
        for (i in 1:I){
            s[i] ~ dbeta(1,1)
            g[i] ~ dbeta(1,1) T(,1-s[i])
        }
    }"""
    jags_ho_dina_model = pyjags.Model(code=jags_hodina_model_str,
                                        data=jags_ho_dina_data,
                                        chains=chains,
                                        threads=cores,
                                        generate_data=True,
                                        chains_per_thread=1,
                                        adapt=False,
                                      progress_bar=True
                                        )
    print("burning.......")
    jags_tarce_prior = jags_ho_dina_model.sample(iterations=burning, vars=[], thin=1)
    print("sampling.......")
    jags_tarce_posterior = jags_ho_dina_model.sample(iterations=iteration, vars=jags_model_return_vars, thin=1)
    x = (jags_tarce_posterior["alpha"].mean(axis=-1).mean(axis=-1) > 0.5)
    if return_trace==True:
        return x,jags_tarce_posterior
    return x


def long_dina_parameter_estimate(T,N,Q,I,K,simu_resp,burning=5000,iteration=5000,chains=2,cores=4
                                 ,return_trace=False,return_full_vars=False):
    jags_data = {
        "Y": simu_resp,
        "Q": Q,
        "N": N,
        "I": I,
        "K": K,
        "T": T
    }
    if return_full_vars==True:
        jags_return_vars = ["mu_theta", "pr_theta", "theta", "att_possi",
                        "att", "theta_beta", "theta_lambda", "lambda_0", "lambda_k", "p"]
    else :
        jags_return_vars = ["att","att_possi"]

    long_dina_jags_model = ling_dina_jags_model = \
        """
        model{
            for (t in 1:T){
                for (n in 1:N){
                    for (i in 1:I){
                        for (k in 1:K){w[t,n,i,k] <- (att[t,n,k]>=Q[t,i,k])}
                        eta[t,n,i] <- prod(w[t,n,i,1:K])
                        logit(p[t,n,i]) <- lambda_0[t,i]+lambda_k[t,i]*eta[t,n,i]
                        Y[t,n,i] ~ dbern(p[t,n,i])
                    }
                }
    
    
                for (n in 1:N){
                    for (k in 1:K){
                        logit(att_possi[t,n,k]) <- theta_beta[k] + theta_lambda[k]*theta[t,n]
                        att[t,n,k] ~ dbern(att_possi[t,n,k])
                    }
                }
            }
    
            for (n in 1:N){
                theta[1:T,n] ~ dmnorm(mu_theta[1:T], pr_theta[1:T,1:T])
            }
            for (k in 1:K){
                theta_beta[k]~ dnorm(0,0.25)
                theta_lambda[k] ~ dnorm(0,0.25) T (0,)
            }
    
            for (t in 1:T){
                for (i in 1:I){
                    lambda_0[t,i] ~ dnorm(-1.096, 0.25)
                    lambda_k[t,i] ~ dnorm(0, 0.25) T(0, )
                }
            }
    
            mu_theta[1] <- 0
            for (t in 2:T){
                mu_theta[t] ~ dnorm(0,0.5)
            }
    
            L_theta[1,1] <- 1
            for (tt in 2:T){
                L_theta[tt,tt] ~ dgamma(1,1)
                for (ttt in 1:(tt-1)){
                    L_theta[tt,ttt] ~ dnorm(0,1)
                    L_theta[ttt,tt] <- 0
                }
            }
            sigma_theta = L_theta %*% t(L_theta)
            pr_theta[1:T, 1:T] <- inverse(sigma_theta[1:T, 1:T])
    
    
        }
        """
    jags_long_dina_model = pyjags.Model(code=long_dina_jags_model,
                                        data=jags_data,
                                        chains=chains,
                                        threads=cores,
                                        generate_data=True,
                                        chains_per_thread=1,
                                        adapt=False,
                                        progress_bar=True
                                        )
    print("burning.......")
    jags_tarce_prior = jags_long_dina_model.sample(iterations=burning, vars=[], thin=1)
    print("sampling.......")
    jags_tarce_posterior = jags_long_dina_model.sample(iterations=iteration, vars=jags_return_vars, thin=1)

    x = (jags_tarce_posterior["att"].mean(axis=-1).mean(axis=-1) > 0.5)
    if return_trace==True:
        return x,jags_tarce_prior
    return x

